from ray.tune.suggest.ax import AxSearch
import numpy as np


class AxSearchExt(AxSearch):

  def __init__(self,
               ax_client,
               mode="max",
               use_early_stopped_trials=None,
               max_concurrent=None,
               save_path=None):
    super(AxSearchExt, self).__init__(ax_client, mode, use_early_stopped_trials,
                                      max_concurrent)
    self.save_path = save_path

  def _process_result(self, trial_id, result):
    ax_trial_index = self._live_trial_mapping[trial_id]
    metric_dict = {self._objective_name: (result[self._objective_name], 0.0)}
    outcome_names = [
        oc.metric.name
        for oc in self._ax.experiment.optimization_config.outcome_constraints
    ]
    metric_dict.update({on: (result[on], 0.0) for on in outcome_names})

    # deal with NaN
    if np.isnan(np.sum(list(metric_dict.values()))):
      self._ax.log_trial_failure(trial_index=ax_trial_index)
    else:
      self._ax.complete_trial(trial_index=ax_trial_index, raw_data=metric_dict)
    self._ax.save_to_json_file(self.save_path)